Repository navigation
perf(ioloader): decode NPY payloads unboxed, and per leading-axis block concurrently - #29
Conversation
07c5540 to
6d824f7
Compare
|
Thanks, Nicolas. We still want to merge this performance work, but the data-loader cleanup in Could you rebase onto the current After the rebase, please rerun |
6d824f7 to
a9d1541
Compare
|
Hey Nicolas, can you rebase this onto the latest main? |
|
Thanks! The NPY decoding improvements still look useful. Could you rebase this onto current main and update the tests for the current loader, preserving its input validation? |
a9d1541 to
181edd4
Compare
|
Rebased onto current What changed in the port:
Re-measured on this tree, same binary for all rows (details in the description):
Checks on The PR description is updated to match. |
Brings in the upstream PR branch (lean-dojo#29) at 181edd4, rebased onto upstream main b062b9a. combined is rebuilt from b062b9a, the first upstream tip on the LibTorch backend (csrc/cuda is gone). Its previous tip d162950 (Lean 4.33, custom CUDA kernels, arena, cache cap, TexTable) is kept as bkp/pre-libtorch/combined and is not merged: nothing on it applies to the LibTorch tree.
|
Hey Nicolas, thanks! This looks like a useful improvement. The loader, its tests, and the updated data API compile locally. My isolated interpreter run could not load the native parseNpy implementation, so I have not confirmed the runtime suite locally yet; that is a scratch-run setup limitation, not evidence of a bug in this PR. Could you share the native runtime test command and results? It would also be good to cover signed zero, infinities, and NaNs explicitly in the decode tests. Keeping this open until the runtime check is confirmed. |
…ck concurrently The `.npy` element loop allocated on every step: an `Option` per byte, an `Except` per element, and a heap cell per element to box the `Float`. It also re-tested the dtype once per element. `decodeNpyPayload` now selects a range decoder once per file and fills a `FloatArray`, so `NpyData.values` is unboxed. `NN.API.Data.Sources` wraps that buffer into a tensor through `Tensor.from` without a copy or a boxed intermediate. Header validation, bounded handle reads, and the truncation checks are unchanged; the in-memory and file-backed readers share one decoder. `parseNpyBlocks` and `readNpyBlocks` return one `FloatArray` per index along the leading axis of a C-order array and decode those blocks concurrently. `Parallelism` and `resolveTaskCount` bound the task count by the request, the block count, and the payload size. `NN/Tests/Data/IO/Npy.lean` checks bit-for-bit round trips for both dtypes, that every task count returns the same blocks, determinism across repeated concurrent decodes, the task-count bounds, Fortran reordering, rejections, and that the file-backed readers agree with the in-memory ones.
… NPY decode path The round-trip tests draw their values from a smooth sample, so none of them is a signed zero, an infinity or a NaN, and `==` on `Float` could not have told the first two apart or accepted the third. `testSpecialValues` stores those bit patterns, with NaN payloads and the extremes of either format, as `<f8` and `<f4` files and checks the flat, block and file-backed decoders. Signed zeros, subnormals, infinities and every finite value decode to exactly their stored bits. A NaN decodes as a NaN, but not as its stored bits: Lean's `Float` is logically `Float.Model`, a subtype of `UInt64` whose `Format.Valid` admits one NaN pattern, and `Float.ofBits` and `Float32.ofBits` return that canonical NaN for every NaN input, sign and payload included, so no decoder over `Float` can preserve them. The test pins the canonical result, and the loader's documentation says so.
181edd4 to
641461c
Compare
|
Rebased onto current The rebase changed two things, and the decoder itself is otherwise line for line what you
Native runtime command and results. The suite is the compiled Signed zero, infinities, NaNs. Writing it found something that belongs in the loader's contract. Signed zeros, subnormals, for w in [0x7ff8000000000000, 0xfff8000000000000, 0x7ff00000deadbeef, 0xfff00000deadbeef] do
let f := Float.ofBits w
IO.println s!"in={w} toBits={f.toBits} isNaN={f.isNaN}"
for w in [0x7fc00000, 0xffc00000, 0x7f80beef, 0xff80beef] do
let g := Float32.ofBits w
IO.println s!"f32 in={w} toBits={g.toBits} widened={g.toFloat.toBits} isNaN={g.isNaN}"So no decoder over |
Body
What
The
.npyelement loop allocates on every step: anOptionper byte insidereadLittleEndian,an
Exceptper element fromreadNpyElement, and a heap cell per element again to store theFloatin anArray Float. It also re-tests the dtype string once per element.This selects the decoder once per file inside
decodeNpyPayload, reads the bytes without anintermediate
Option, and decodes into aFloatArray.NpyData.valuesbecomes aFloatArray:the payload is bulk numeric data, so it costs one machine word per element instead of a word
plus a heap cell, and
Tensor.fromwraps that buffer into a tensor without copying it. Headervalidation, the bounded handle reads, prefix validation, and the truncation checks are unchanged;
the in-memory and file-backed readers share the one decoder.
It also adds
parseNpyBlocksandreadNpyBlocks, which return oneFloatArrayper index along the leading axis. Those blocks are physically contiguous in aC-order payload, so they are independent and can be decoded concurrently. Callers that want the
blocks apart anyway (one channel, one band, one batch element) get the concurrency for free.
Callers that want a single flat payload keep using
parseNpy/readNpy, which stay sequential:producing one contiguous result concurrently would mean filling several arrays and then copying
them all into one, and that copy walks every element a second time for several times the cost of
the concurrency it was meant to buy.
ParallelismandresolveTaskCountmake the task count explicit rather than tying it to thefile's shape. The count follows the machine and the block count follows the file, and those are
independent, so blocks are grouped across tasks rather than handed out one per task: a
(50000, 1, 28, 28)dataset does not spawn 50 000 tasks.Measured
Measured on the Lean v4.34.0 tree, before this rebase onto 4.35
main; the decoder is unchanged bythe rebase apart from the names of the block readers. One 764 MB
<f8file of shape(8, 3456, 3456)(95 551 488 elements), already in memory, Xeon w7-3465X (28 cores, 56 threads).One binary holds all three decoders; the first row is the element loop
mainships atb062b9a3,copied verbatim. Five trials per row, each in its own process so peak RSS is per path; the rate is
file bytes over the median; every trial passes a distinct tag so the pure decode cannot be shared
across trials. Peak RSS includes the 764 MB input held in memory.
mainparseNpy, unboxed sequential.sequential.tasks 8.auto.autoresolves to 8 tasks here: the leading dimension is 8, andresolveTaskCountneversplits a block, so the file's leading axis is the ceiling on this decode's parallelism whatever
the core count.
Tests
NN/Tests/Data/IO/Npy.leansynthesises.npyfiles in memory, so the expected values areknown exactly, and writes the same bytes to
.lake/build/tmpfor the file-backed readers.It checks:
<f8and<f4, scalars included.==onFloatwould be the wrong test: it identifies0.0with-0.0and makes everyNaNunequal to itself, so it would both miss sign-bit transcription errors and reject acorrectly decoded
NaN;==onFloatcannot see, through every path (parseNpy, the blocks,readNpy,readNpyBlocks): both signed zeros, both infinities and the extremes of eitherformat decode to exactly their stored bits; a NaN decodes as the canonical quiet NaN whatever
its sign or payload, because Lean's
Float.ofBitsandFloat32.ofBitscanonicalize every NaNpattern (
Float.Model.Format.Validadmits one NaN encoding). The test pins that result and theloader's docstring states it;
parseNpyreturns, for 1-D, 2-D,and 3-D shapes and for zero-sized axes on either side;
above the block count and above the core count, plus
sequentialandauto. This is whatmakes the task count a free tuning knob rather than part of the contract;
resolveTaskCountrespects each of its three bounds;dtype, Fortran order and scalar shape for the block path;
readNpyandreadNpyBlocksagree bit for bit with their in-memory counterparts,and the file-backed block reader rejects a Fortran-order or scalar header without reading a
payload (the fixture is a bare header, and the error names the storage order, not truncation).
It is the only NPY test in the curated suite:
mainno longer hasNN/Tests/API/Data.lean.Note on
Task.spawnThe block decode uses
Task.spawn, which takes the work as aUnit → αclosure. That isload-bearing and there is a comment saying so: handing a pure expression to an
IO-level spawninstead lets the compiler float it out of the closure onto the spawning thread, and the fan-out
then runs sequentially at exactly the single-threaded rate while still looking parallel.
Compatibility
NpyData.valueschanges type fromArray FloattoFloatArray. In-tree consumer:NN/API/Data/Sources.lean(nowTensor.fromon the buffer).NN/Examples/Data/Loaders/Npy.leanreads only
dtypeandshape.Checks run
On
641461c1(two commits onmain72754452, Lean v4.35.0), compiled suite through the testdriver: